# Configuration - CHANGE THIS TO YOUR EXPERIMENT FOLDER
EXP_PATH = "/Users/ettoremiglioranza/Projects/NewAM/data/experiments/EXP_20260131_182446_128x64x64" 

# Backend for PyVista in Jupyter
import pyvista as pv
pv.set_jupyter_backend('static')  # Use 'static' for screenshots, 'ipyvtklink' for interactive

from topopt_ml.io import load_sample, load_dataset_index

# Load dataset index

dataset_index = load_dataset_index(EXP_PATH)
print(f"Found {len(dataset_index)} samples in dataset")

print("\nSample metadata:")
for i, sample in enumerate(dataset_index[:3]):
    print(f"  Sample {sample['sample_id']}: Load at {sample['load_center']}, "
          f"Radius: {sample['load_radius']:.2f}, Time: {sample['solve_time']:.2f}s")

Found 10 samples in dataset

Sample metadata:
  Sample 0001: Load at [64, 34], Radius: 7.13, Time: 158.66s
  Sample 0002: Load at [48, 17], Radius: 7.89, Time: 129.71s
  Sample 0003: Load at [39, 34], Radius: 6.87, Time: 50.15s

# Load first sample
sample_id = dataset_index[0]['sample_id']
X, Y = load_sample(EXP_PATH, sample_id)

print(f"Input tensor shape: {X.shape}")
print(f"Target tensor shape: {Y.shape}")
print(f"\nInput channels:")
print(f"  Ch 0 (Solid): {X[:,:,:,0].sum():.0f} elements")
print(f"  Ch 3 (Fz): {X[:,:,:,3].sum():.2f} N")
print(f"\nDensity statistics:")
print(f"  Range: [{Y.min():.4f}, {Y.max():.4f}]")
print(f"  Mean: {Y.mean():.4f}")

Input tensor shape: (128, 64, 64, 4)
Target tensor shape: (128, 64, 64)

Input channels:
  Ch 0 (Solid): 524288 elements
  Ch 3 (Fz): -2000.00 N

Density statistics:
  Range: [0.0000, 1.0000]
  Mean: 0.1500

from topopt_ml.io import iterate_samples

# Validate multiple samples
for sid, X, Y in iterate_samples(EXP_PATH, max_samples=15):
    print(f"\n{'='*60}")
    print(f"Sample: {sid}")
    print(f"{'='*60}")
    validate_sample_visual(X, Y, sample_id=sid, threshold=0.3, show_plot=True)

============================================================
Sample: 0001
============================================================

  Physical Consistency Checks:
    ✓ Load applied: True (-2000.00 N)
    ✓ Structure exists: True
    ✓ Volume fraction: 0.150
    ✓ Load center: (63.5, 33.5, 63.0)
    ✓ Structure-load distance: 37.7 elements

Sample 0001:
  Input shape: (128, 64, 64, 4)
  Target shape: (128, 64, 64)
  Force sum (Fz): -2000.00 N
  Density range: [0.0000, 1.0000]
  Density mean: 0.1500

  ✅ Sample 0001 validated


============================================================
Sample: 0002
============================================================

  Physical Consistency Checks:
    ✓ Load applied: True (-2000.00 N)
    ✓ Structure exists: True
    ✓ Volume fraction: 0.150
    ✓ Load center: (47.5, 16.5, 63.0)
    ✓ Structure-load distance: 49.2 elements

Sample 0002:
  Input shape: (128, 64, 64, 4)
  Target shape: (128, 64, 64)
  Force sum (Fz): -2000.00 N
  Density range: [0.0000, 1.0000]
  Density mean: 0.1500

  ✅ Sample 0002 validated


============================================================
Sample: 0003
============================================================

  Physical Consistency Checks:
    ✓ Load applied: True (-2000.00 N)
    ✓ Structure exists: True
    ✓ Volume fraction: 0.150
    ✓ Load center: (38.5, 33.5, 63.0)
    ✓ Structure-load distance: 31.1 elements

Sample 0003:
  Input shape: (128, 64, 64, 4)
  Target shape: (128, 64, 64)
  Force sum (Fz): -2000.00 N
  Density range: [0.0000, 1.0000]
  Density mean: 0.1500

  ✅ Sample 0003 validated


============================================================
Sample: 0004
============================================================

  Physical Consistency Checks:
    ✓ Load applied: True (-2000.00 N)
    ✓ Structure exists: True
    ✓ Volume fraction: 0.150
    ✓ Load center: (52.5, 37.5, 63.0)
    ✓ Structure-load distance: 47.0 elements

Sample 0004:
  Input shape: (128, 64, 64, 4)
  Target shape: (128, 64, 64)
  Force sum (Fz): -2000.00 N
  Density range: [0.0000, 1.0000]
  Density mean: 0.1500

  ✅ Sample 0004 validated


============================================================
Sample: 0005
============================================================

  Physical Consistency Checks:
    ✓ Load applied: True (-2000.00 N)
    ✓ Structure exists: True
    ✓ Volume fraction: 0.150
    ✓ Load center: (62.5, 22.5, 63.0)
    ✓ Structure-load distance: 41.9 elements

Sample 0005:
  Input shape: (128, 64, 64, 4)
  Target shape: (128, 64, 64)
  Force sum (Fz): -2000.00 N
  Density range: [0.0000, 1.0000]
  Density mean: 0.1500

  ✅ Sample 0005 validated


============================================================
Sample: 0006
============================================================

  Physical Consistency Checks:
    ✓ Load applied: True (-2000.00 N)
    ✓ Structure exists: True
    ✓ Volume fraction: 0.150
    ✓ Load center: (35.5, 46.5, 63.0)
    ✓ Structure-load distance: 30.5 elements

Sample 0006:
  Input shape: (128, 64, 64, 4)
  Target shape: (128, 64, 64)
  Force sum (Fz): -2000.00 N
  Density range: [0.0000, 1.0000]
  Density mean: 0.1500

  ✅ Sample 0006 validated


============================================================
Sample: 0007
============================================================

  Physical Consistency Checks:
    ✓ Load applied: True (-2000.00 N)
    ✓ Structure exists: True
    ✓ Volume fraction: 0.150
    ✓ Load center: (37.5, 22.5, 63.0)
    ✓ Structure-load distance: 31.0 elements

Sample 0007:
  Input shape: (128, 64, 64, 4)
  Target shape: (128, 64, 64)
  Force sum (Fz): -2000.00 N
  Density range: [0.0000, 1.0000]
  Density mean: 0.1500

  ✅ Sample 0007 validated


============================================================
Sample: 0008
============================================================

  Physical Consistency Checks:
    ✓ Load applied: True (-2000.00 N)
    ✓ Structure exists: True
    ✓ Volume fraction: 0.150
    ✓ Load center: (78.5, 16.5, 63.0)
    ✓ Structure-load distance: 38.5 elements

Sample 0008:
  Input shape: (128, 64, 64, 4)
  Target shape: (128, 64, 64)
  Force sum (Fz): -2000.00 N
  Density range: [0.0000, 1.0000]
  Density mean: 0.1500

  ✅ Sample 0008 validated


============================================================
Sample: 0009
============================================================

  Physical Consistency Checks:
    ✓ Load applied: True (-2000.00 N)
    ✓ Structure exists: True
    ✓ Volume fraction: 0.150
    ✓ Load center: (89.5, 29.5, 63.0)
    ✓ Structure-load distance: 40.8 elements

Sample 0009:
  Input shape: (128, 64, 64, 4)
  Target shape: (128, 64, 64)
  Force sum (Fz): -2000.00 N
  Density range: [0.0000, 1.0000]
  Density mean: 0.1500

  ✅ Sample 0009 validated


============================================================
Sample: 0010
============================================================

  Physical Consistency Checks:
    ✓ Load applied: True (-2000.00 N)
    ✓ Structure exists: True
    ✓ Volume fraction: 0.150
    ✓ Load center: (68.5, 19.5, 63.0)
    ✓ Structure-load distance: 37.9 elements

Sample 0010:
  Input shape: (128, 64, 64, 4)
  Target shape: (128, 64, 64)
  Force sum (Fz): -2000.00 N
  Density range: [0.0000, 1.0000]
  Density mean: 0.1500

# 8. Structural Diversity Analysis

import numpy as np
import matplotlib.pyplot as plt
from itertools import combinations

print("=" * 60)
print("STRUCTURAL DIVERSITY ANALYSIS")
print("=" * 60)

# Load all samples
all_densities = []
sample_ids = []
for sample in dataset_index:
    sid = sample['sample_id']
    _, Y = load_sample(EXP_PATH, sid)
    all_densities.append(Y)
    sample_ids.append(sid)

print(f"\nAnalyzing {len(all_densities)} samples...")

# Compute pairwise Jaccard similarities
threshold = 0.3
n_samples = len(all_densities)
similarity_matrix = np.zeros((n_samples, n_samples))

print(f"\nPairwise Jaccard Similarities (threshold={threshold}):")
print("-" * 60)

for i in range(n_samples):
    for j in range(n_samples):
        if i == j:
            similarity_matrix[i, j] = 1.0
        elif i < j:
            binary1 = (all_densities[i] > threshold).astype(float)
            binary2 = (all_densities[j] > threshold).astype(float)
            
            intersection = (binary1 * binary2).sum()
            union = ((binary1 + binary2) > 0).sum()
            jaccard = intersection / union if union > 0 else 0
            
            similarity_matrix[i, j] = jaccard
            similarity_matrix[j, i] = jaccard
            
            print(f"  Sample {sample_ids[i]} ↔ {sample_ids[j]}: {jaccard:.4f}")

# Get upper triangle (excluding diagonal) for statistics
upper_triangle = similarity_matrix[np.triu_indices(n_samples, k=1)]

print("\n" + "=" * 60)
print("DIVERSITY STATISTICS")
print("=" * 60)
print(f"Mean pairwise similarity: {np.mean(upper_triangle):.4f}")
print(f"Std pairwise similarity:  {np.std(upper_triangle):.4f}")
print(f"Min similarity:           {np.min(upper_triangle):.4f}")
print(f"Max similarity:           {np.max(upper_triangle):.4f}")

# Interpretation
mean_sim = np.mean(upper_triangle)
print("\n" + "-" * 60)
print("INTERPRETATION:")
print("-" * 60)
if mean_sim > 0.7:
    print("⚠️  HIGH SIMILARITY (>0.7): Poor diversity - samples too similar!")
    print("   Consider: more varied load positions, different BCs, or")
    print("   multiple load cases per sample.")
elif mean_sim > 0.3:
    print("⚠️  MODERATE SIMILARITY (0.3-0.7): Acceptable but could improve")
    print("   Samples show some variation but may benefit from more diversity.")
else:
    print("✅ GOOD DIVERSITY (<0.3): Structures are sufficiently different")
    print("   Dataset has healthy variation for ML training.")

# Visualization
fig = plt.figure(figsize=(16, 6))

# Left subplot: Similarity heatmap
ax1 = plt.subplot(1, 2, 1)
im = ax1.imshow(similarity_matrix, cmap='RdYlGn_r', vmin=0, vmax=1)
ax1.set_xticks(range(n_samples))
ax1.set_yticks(range(n_samples))
ax1.set_xticklabels(sample_ids)
ax1.set_yticklabels(sample_ids)
ax1.set_xlabel('Sample ID')
ax1.set_ylabel('Sample ID')
ax1.set_title('Pairwise Jaccard Similarity Matrix\n(Lower = More Different)')
plt.colorbar(im, ax=ax1, label='Similarity')

# Add text annotations
for i in range(n_samples):
    for j in range(n_samples):
        text = ax1.text(j, i, f'{similarity_matrix[i, j]:.2f}',
                       ha="center", va="center", 
                       color="white" if similarity_matrix[i, j] > 0.5 else "black",
                       fontsize=9)

# Right subplot: Visual comparison of mid-plane slices
ax2 = plt.subplot(1, 2, 2)
ax2.axis('off')

# Create grid of subplots for slices
n_cols = n_samples
slice_height = 1.0 / n_cols

for idx, (Y, sid) in enumerate(zip(all_densities, sample_ids)):
    # Create inset axes for each slice
    left = 0.05
    bottom = 0.95 - (idx + 1) * slice_height + 0.02
    width = 0.9
    height = slice_height - 0.04
    
    ax_slice = fig.add_axes([0.52 + left * 0.45, bottom, width * 0.45, height])
    
    # Show mid-plane (x=64 for 128-element domain)
    mid_slice = Y[64, :, :]
    ax_slice.imshow(mid_slice.T, cmap='gray', vmin=0, vmax=1, origin='lower', aspect='auto')
    ax_slice.set_ylabel(f'{sid}', fontsize=9)
    ax_slice.set_xticks([])
    ax_slice.set_yticks([])
    
    if idx == 0:
        ax_slice.set_title('Mid-Plane Slices (X=64)', fontsize=10)

plt.suptitle(f'Structural Diversity Analysis - {len(all_densities)} Samples', 
             fontsize=14, fontweight='bold', y=0.98)

# Save figure
import os
# Save figure
save_path = 'structural_diversity.png'  # Save in current directory
plt.savefig(save_path, dpi=150, bbox_inches='tight')
print(f"\n📊 Diversity visualization saved to: {save_path}")

plt.show()

print("\n" + "=" * 60)
print("If slices look substantially different → Good diversity ✓")
print("If slices look nearly identical → Need more varied BCs")
print("=" * 60)

============================================================
STRUCTURAL DIVERSITY ANALYSIS
============================================================

Analyzing 10 samples...

Pairwise Jaccard Similarities (threshold=0.3):
------------------------------------------------------------
  Sample 0001 ↔ 0002: 0.0370
  Sample 0001 ↔ 0003: 0.1252
  Sample 0001 ↔ 0004: 0.1003
  Sample 0001 ↔ 0005: 0.0769
  Sample 0001 ↔ 0006: 0.1071
  Sample 0001 ↔ 0007: 0.1076
  Sample 0001 ↔ 0008: 0.1985
  Sample 0001 ↔ 0009: 0.3830
  Sample 0001 ↔ 0010: 0.2961
  Sample 0002 ↔ 0003: 0.0080
  Sample 0002 ↔ 0004: 0.1773
  Sample 0002 ↔ 0005: 0.5089
  Sample 0002 ↔ 0006: 0.0000
  Sample 0002 ↔ 0007: 0.0215
  Sample 0002 ↔ 0008: 0.2318
  Sample 0002 ↔ 0009: 0.0776
  Sample 0002 ↔ 0010: 0.1988
  Sample 0003 ↔ 0004: 0.0076
  Sample 0003 ↔ 0005: 0.0000
  Sample 0003 ↔ 0006: 0.5494
  Sample 0003 ↔ 0007: 0.5885
  Sample 0003 ↔ 0008: 0.0853
  Sample 0003 ↔ 0009: 0.0783
  Sample 0003 ↔ 0010: 0.0968
  Sample 0004 ↔ 0005: 0.3391
  Sample 0004 ↔ 0006: 0.0024
  Sample 0004 ↔ 0007: 0.0009
  Sample 0004 ↔ 0008: 0.0574
  Sample 0004 ↔ 0009: 0.0859
  Sample 0004 ↔ 0010: 0.0686
  Sample 0005 ↔ 0006: 0.0000
  Sample 0005 ↔ 0007: 0.0000
  Sample 0005 ↔ 0008: 0.1678
  Sample 0005 ↔ 0009: 0.0987
  Sample 0005 ↔ 0010: 0.1552
  Sample 0006 ↔ 0007: 0.3158
  Sample 0006 ↔ 0008: 0.0340
  Sample 0006 ↔ 0009: 0.0510
  Sample 0006 ↔ 0010: 0.0469
  Sample 0007 ↔ 0008: 0.1222
  Sample 0007 ↔ 0009: 0.0797
  Sample 0007 ↔ 0010: 0.1222
  Sample 0008 ↔ 0009: 0.3032
  Sample 0008 ↔ 0010: 0.6071
  Sample 0009 ↔ 0010: 0.3577

============================================================
DIVERSITY STATISTICS
============================================================
Mean pairwise similarity: 0.1573
Std pairwise similarity:  0.1633
Min similarity:           0.0000
Max similarity:           0.6071

------------------------------------------------------------
INTERPRETATION:
------------------------------------------------------------
✅ GOOD DIVERSITY (<0.3): Structures are sufficiently different
   Dataset has healthy variation for ML training.

📊 Diversity visualization saved to: structural_diversity.png

============================================================
If slices look substantially different → Good diversity ✓
If slices look nearly identical → Need more varied BCs
============================================================

Topology Optimization Dataset - Visualization & Validation¶

Load Dataset¶

Multi-Sample Gallery¶